package com.example.langchain4j.controller;

import com.example.langchain4j.service.NewsService;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.SystemMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import dev.langchain4j.model.ollama.OllamaEmbeddingModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.Arrays;
import java.util.List;
import java.util.stream.Collectors;

@RestController
public class OllamaController {

    interface Assistant {
        Flux<String> chat(String userMessage);
    }

    @Autowired
    private OllamaEmbeddingModel embeddingModel;

    @Autowired
    private ChatLanguageModel chatModel;

    @Autowired
    private StreamingChatLanguageModel streamingChatLanguageModel;

    @GetMapping("/ai")
    String generation(String userInput) {
        return chatModel.chat(new SystemMessage("请全部用简体中文回答"), new UserMessage(userInput)).aiMessage().text();
    }

    @GetMapping(value = "/streamAi", produces = "text/html;charset=UTF-8")
    Flux<String> streamGeneration(@RequestParam(value = "userInput", defaultValue = "给我讲一个笑话") String userInput) {
        return Flux.create(emitter -> {
            List<ChatMessage> chatMessageList = Arrays.asList(new SystemMessage("请全部用简体中文回答"), new UserMessage(userInput));
            streamingChatLanguageModel.chat(chatMessageList, new StreamingChatResponseHandler() {
                @Override
                public void onPartialResponse(String s) {
                    emitter.next(s);
                }

                @Override
                public void onCompleteResponse(ChatResponse chatResponse) {
                    emitter.complete();
                }

                @Override
                public void onError(Throwable throwable) {
                    emitter.error(throwable);
                }
            });
        });
    }

    @GetMapping(value = "/toolCallGeneration", produces = "text/html;charset=UTF-8")
    Flux<String> toolCallGeneration(@RequestParam(value = "userInput", defaultValue = "给我讲一个笑话") String userInput) {
        return AiServices.builder(Assistant.class).streamingChatLanguageModel(streamingChatLanguageModel).tools(new NewsService()).build().chat(userInput).onErrorResume(e -> Flux.just("系统繁忙，请稍后重试"));
    }

    @GetMapping(value = "/ragGeneration", produces = "text/html;charset=UTF-8")
    Flux<String> ragGeneration(@RequestParam(value = "userInput", defaultValue = "synchronized和volatile有什么区别？") String userInput) {
        //从本地txt文件读取文本内容
        Document document = FileSystemDocumentLoader.loadDocument("C:\\Users\\18716\\Desktop\\JAVA-RAG.txt", new TextDocumentParser());
        //文本内容分割
        DocumentSplitter splitter = DocumentSplitters.recursive(220, 0);
        List<TextSegment> segments = splitter.split(document);
        //文本内容向量化并存储
        List<Embedding> embeddings = embeddingModel.embedAll(segments).content();
        InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
        embeddingStore.addAll(embeddings, segments);
        ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .maxResults(5)
                .minScore(0.6)
                .build();
        return AiServices.builder(Assistant.class).streamingChatLanguageModel(streamingChatLanguageModel).contentRetriever(contentRetriever).build().chat(userInput).onErrorResume(e -> Flux.just("系统繁忙，请稍后重试"));
    }

}
